'''
作为预测代码，可用于对两个个模型的预测
1. manifold_3d_train.py: manifold_3d_*
2. remix_3d_train.py: remix_3d_*
'''

import os
import pandas as pd
import numpy as np

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from net_two_manifold_mixup_remix_3d import generate_model
from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc



BatchSize = 16
save_path = '/home/zmy/pytorch_code/checkpoint/two_manifold_3d_1.2_20210523_132741'
weight_name = 'two_manifold_3d_1.2-99-0.5618634715676307-regular.pth'

weights_path = save_path+'/'+weight_name

# 自定义数据集和数据预处理
class MyDataset(Dataset):

    def __init__(self, datalist):
        self.data_info = datalist


    def __len__(self):
        return len(self.data_info)

    def __getitem__(self, item):
        patientID = self.data_info[item][0]
        label = self.data_info[item][1]
        img_path = self.data_info[item][2]

        img = np.load(img_path)

        # name_list = os.listdir(ct_path)
        # name_list.sort()
        #
        # ct_3D_array = []
        # pet_3D_array = []
        #
        # for i in range(len(name_list)):
        #     ct = np.load(ct_path+name_list[i])
        #     pet = np.load(pet_path+name_list[i])
        #
        #     ct_3D_array.append(ct)
        #     pet_3D_array.append(pet)
        #
        #
        # # ct和pet进行合并
        # img = np.asarray([ct_3D_array, pet_3D_array], dtype=np.float)
        #
        # # print('img shape: ', img.shape)
        #
        # # 归一化输入大小
        # resize_para = 8.0 / len(name_list)
        # img = scipy.ndimage.zoom(img, [1, resize_para, 0.5, 0.5])
        #
        # # print('resize img shape: ', img.shape)


        return {'image': torch.from_numpy(img), 'label': torch.tensor(label)}


# 读取文件列表
def read_csv(data_sets):
    sets_path = '/data1/zmy/data2021/origin_data/divide_csv/five/'

    # 读取数据集
    data_features = []

    for set in data_sets:
        train_data = pd.read_csv(sets_path+set)
        for j in range(len(train_data)):

            # 读取文件地址
            patientid = train_data['patientID'][j]

            img_path = '/data1/zmy/data2021/auge2_data/img_3D/'+str(patientid)+'/img3d.npy'

            cancer_type = int(train_data['cancer_type'][j])-1
            if cancer_type==0 or cancer_type==1:
                one_feature = [patientid, cancer_type, img_path]

                data_features.append(one_feature)


    return data_features



# 创建网络
def resnet():

    # 创建resnet50网络
    net = generate_model(50)

    # 打印网络结构和参数量
    # print(net)
    # print("Total number of paramerters in networks is {}  ".format(sum(x.numel() for x in net.parameters())))

    return net




# 结果预测
def predict():
    device = torch.device("cuda:3")

    # 创建网络结构
    net = resnet().to(device)

    # 加载测试数据
    test_data_sets = ['test.csv']
    test_list = read_csv(test_data_sets)
    test_dataset = MyDataset(test_list)
    testloader = DataLoader(test_dataset, batch_size=BatchSize, shuffle=False, num_workers=2)

    # 加载模型权重
    net.load_state_dict(torch.load(weights_path))
    net.eval()

    # 返回的预测结果及对应标签
    y_preds = []

    y_preds_ones = []  # 预测为腺癌的概率

    y_true = []

    with torch.no_grad():
        for n_iter, data in enumerate(testloader):
            # print('iteration:{}\ttotal {} iterations'.format(n_iter + 1, len(testloader)))

            images = data['image'].type(torch.FloatTensor).to(device)
            labels = data['label'].to(device)

            outputs = net(images, labels, use_mixup=False, mixup_alpha=0.1, layer_mix=None)
            _, preds = outputs.max(1)

            outputs = torch.softmax(outputs, dim=1)
            result = outputs.tolist()
            for mm in range(len(result)):
                y_preds_ones.append(result[mm][1])

            y_preds.extend(preds.tolist())
            y_true.extend(labels.tolist())

            # print(preds.tolist())
            # print(labels.tolist())
        print(y_preds_ones)
        print(y_true)
        # 保存测试集结果（概率值）
        np.save('/home/zmy/pytorch_code/two_result/two_manifold_labels.npy', y_true)
        np.save('/home/zmy/pytorch_code/two_result/two_manifold_preds.npy', y_preds_ones)
    return y_preds, y_preds_ones, y_true


# 结果评估
def evaluation(preds, y_preds_one, labels):

    # 计算评估指标
    target_names = ['1', '2']
    result_statis = classification_report(y_true=labels, y_pred=preds, target_names=target_names)
    print(result_statis)

    # 计算混淆矩阵
    confusion = confusion_matrix(y_true=labels, y_pred=preds)
    print(confusion)

    # 计算auc
    fpr, tpr, threshold = roc_curve(y_true, y_preds_one)
    roc_auc = auc(fpr, tpr)
    print('roc is: ', roc_auc)


if __name__ == '__main__':

    # 预测结果
    y_preds, y_preds_one, y_true = predict()

    print('weight name: ', weight_name)

    # 结果评估
    evaluation(y_preds, y_preds_one, y_true)


